from typing import TypedDict
from langgraph.graph import StateGraph
import requests
from datetime import datetime
from novel_prompt import refine_outline_pt, outline_pt, chapter_pt


class NovelNodes:
    def __init__(self):
        self.endpoint = 'https://open.bigmodel.cn/api/paas/v4/chat/completions'
        self.api_key = '734c6ee4a5524052a1fc44580660dc38.F5A02rjHWxocAsaT'
        self.model = 'glm-4-flash'

    def _call_llm(self, prompt: str, caller: str) -> str:
        headers = {
            'Authorization': f'Bearer {self.api_key}',
            'Content-Type': 'application/json'
        }
        
        pt_out = prompt[:50]
        pt_out = prompt
        print(f"[{datetime.now()}] {caller}调用LLM - 提示词:\n{pt_out}")
        
        response = requests.post(
            self.endpoint,
            headers=headers,
            json={
                "model": self.model,
                "messages": [{"role": "user", "content": prompt}]
            }
        )
        
        return response.json()['choices'][0]['message']['content']

    def generate_initial_outline(self, state: dict) -> dict:
        prompt = outline_pt + """

# 参考资料
'''
""" + state['ref'] + "'''"

        state['outline'] = self._call_llm(prompt, '生成初始大纲')
        return state

    def refine_outline(self, state: dict) -> dict:
        prompt = f"""
# 当前大纲
'''
{state['outline']}
'''

""" + refine_outline_pt
        state['outline'] = self._call_llm(prompt, '迭代大纲')

        # 保存最终大纲
        with open('output/novel_outline.md', 'w', encoding='utf-8') as f:
            f.write(state['outline'])

        return state

    def write_chapter(self, state: dict) -> dict:
        chapter_prompt = f"""根据大纲创作章节内容
# Task
创作第{state['current_chapter']}章
""" + chapter_pt +"""
# 大纲
'''
""" + state['outline'] 
        
        content = self._call_llm(chapter_prompt, '撰写章节')
        print(f"生成章节内容（前50字）：{content[:50]}")
        
        # 保存章节
        from util import writeFile
        writeFile(f"output/chapter_{state['current_chapter']}.txt", content)
        
        state['chapter_content'] = content
        state['word_count'] = len(content)
        state['current_chapter'] = state['current_chapter'] + 1
        
        return state

    def final_review(self, state: dict) -> dict:
        prompt = f"最终审核完整小说内容：\n{state['chapter_content'][:50]}\n请检查：\n1. 逻辑一致性\n2. 文学质量\n3. 语法错误"
        state['review'] = self._call_llm(prompt, '最终审核')

        return state